#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os, time, traceback
import numpy as np
import rospy
import cv2
from cv_bridge import CvBridge, CvBridgeError
from ultralytics import YOLO
from sensor_msgs.msg import Image, CompressedImage, PointCloud2
from nav_msgs.msg import Odometry
from std_msgs.msg import Header
from sensor_msgs import point_cloud2 as pc2
import message_filters
from scipy.spatial.transform import Rotation as R

def _abs_topic(t): return t if t.startswith("/") else "/" + t

def make_SE3(Rm, tvec):
    T = np.eye(4); T[:3,:3] = Rm; T[:3,3] = tvec.reshape(3); return T

def inv_SE3(T):
    Rm = T[:3,:3]; t = T[:3,3]
    Ti = np.eye(4); Ti[:3,:3] = Rm.T; Ti[:3,3] = -Rm.T @ t; return Ti

class Node(object):
    def __init__(self):
        # ---------- 参数 ----------
        self.img_topic   = _abs_topic(rospy.get_param("~img_topic", "/camera/color/image_raw"))
        self.cloud_topic = _abs_topic(rospy.get_param("~cloud_topic", "/cloud_registered"))
        self.odom_topic  = _abs_topic(rospy.get_param("~odom_topic",  "/aft_mapped_to_init"))

        self.fx = float(rospy.get_param("~cam_fx", 0.0))
        self.fy = float(rospy.get_param("~cam_fy", 0.0))
        self.cx = float(rospy.get_param("~cam_cx", 0.0))
        self.cy = float(rospy.get_param("~cam_cy", 0.0))
        self.cam_w = int(rospy.get_param("~cam_width", 0))
        self.cam_h = int(rospy.get_param("~cam_height", 0))

        try:
            Rcl = np.array(rospy.get_param("~Rcl"), dtype=np.float64).reshape(3,3)
            tcl = np.array(rospy.get_param("~Pcl"), dtype=np.float64).reshape(3,1)
        except Exception as e:
            rospy.logerr("Param Rcl/Pcl missing or bad shape: %s", e); raise
        self.T_l_c = make_SE3(Rcl, tcl)     # LiDAR->Cam

        try:
            Rex = np.array(rospy.get_param("~extrinsic_R"), dtype=np.float64).reshape(3,3)
            tex = np.array(rospy.get_param("~extrinsic_T"), dtype=np.float64).reshape(3,1)
        except Exception as e:
            rospy.logerr("Param extrinsic_R/T missing or bad shape: %s", e); raise
        # 传入的是 IMU <- LiDAR，取逆得 LiDAR -> IMU
        R_li = Rex.T
        t_li = -Rex.T @ tex
        self.T_l_i = make_SE3(R_li, t_li)

        self.model_path = rospy.get_param("~yolo_model_path", "yolov8n.pt")
        self.conf_thres = float(rospy.get_param("~conf_thres", 0.25))
        self.iou_thres  = float(rospy.get_param("~iou_thres", 0.45))
        self.img_size   = int(rospy.get_param("~img_size", 640))
        self.device     = str(rospy.get_param("~device", "cpu"))

        self.min_depth  = float(rospy.get_param("~min_depth", 0.1))
        self.max_depth  = float(rospy.get_param("~max_depth", 80.0))

        # ---------- 打印全部参数 ----------
        rospy.loginfo("===== YOLO-FUSION PARAMS =====")
        rospy.loginfo("img_topic      : %s", self.img_topic)
        rospy.loginfo("cloud_topic    : %s", self.cloud_topic)
        rospy.loginfo("odom_topic     : %s", self.odom_topic)
        rospy.loginfo("intrinsics     : fx=%.3f fy=%.3f cx=%.3f cy=%.3f size=%dx%d",
                      self.fx, self.fy, self.cx, self.cy, self.cam_w, self.cam_h)
        rospy.loginfo("Rcl (LiDAR->Cam):\n%s", Rcl); rospy.loginfo("tcl:\n%s", tcl.reshape(-1))
        rospy.loginfo("extrinsic IMU<-LiDAR (R):\n%s", Rex); rospy.loginfo("t:\n%s", tex.reshape(-1))
        rospy.loginfo("T_l_i (LiDAR->IMU):\n%s", self.T_l_i)
        rospy.loginfo("YOLO: model=%s, img_size=%d, conf=%.2f, iou=%.2f, device=%s",
                      self.model_path, self.img_size, self.conf_thres, self.iou_thres, self.device)
        rospy.loginfo("depth range    : [%.2f, %.2f]", self.min_depth, self.max_depth)

        # ---------- 订阅 ----------
        self.bridge = CvBridge()
        self._image_is_compressed = self.img_topic.endswith("/compressed")
        if self._image_is_compressed:
            img_sub = message_filters.Subscriber(self.img_topic, CompressedImage, queue_size=10, buff_size=2**24)
        else:
            img_sub = message_filters.Subscriber(self.img_topic, Image, queue_size=10, buff_size=2**24)
        cloud_sub = message_filters.Subscriber(self.cloud_topic, PointCloud2, queue_size=10, buff_size=2**24)
        odom_sub  = message_filters.Subscriber(self.odom_topic,  Odometry,    queue_size=30)

        self.pub_cloud = rospy.Publisher("/yolo_fusion/cloud_highlight", PointCloud2, queue_size=1)
        self.pub_img   = rospy.Publisher("/yolo_fusion/image_marked",   Image,       queue_size=1)

        self.msg_cnt_img = 0; self.msg_cnt_cloud = 0; self.msg_cnt_odom = 0
        self.last_stamp_img = None; self.last_stamp_cloud = None; self.last_stamp_odom = None

        ats = message_filters.ApproximateTimeSynchronizer([img_sub, cloud_sub, odom_sub], queue_size=50, slop=0.5)
        ats.registerCallback(self.cb)

        # 周期性自检：长时间未收到消息报警
        rospy.Timer(rospy.Duration(3.0), self.watchdog)

        # ---------- YOLO ----------
        rospy.loginfo("[YOLO] loading %s", self.model_path)
        self.model = YOLO(self.model_path)
        try:
            self.model.fuse()
        except Exception as e:
            rospy.logwarn("model.fuse() failed: %s", e)
        rospy.loginfo("Node ready. Subscribed & waiting...")

    def watchdog(self, _evt):
        now = rospy.Time.now()
        def age(ts): 
            if ts is None: return None
            return (now - ts).to_sec()
        a_img   = age(self.last_stamp_img)
        a_cloud = age(self.last_stamp_cloud)
        a_odom  = age(self.last_stamp_odom)
        if a_img is None or a_img > 3.0:
            rospy.logwarn("No IMAGE for %.1fs. Check topic: %s (compressed? namespace?)", a_img if a_img else -1, self.img_topic)
        if a_cloud is None or a_cloud > 3.0:
            rospy.logwarn("No CLOUD for %.1fs. Check topic: %s (type PointXYZRGB?)", a_cloud if a_cloud else -1, self.cloud_topic)
        if a_odom is None or a_odom > 3.0:
            rospy.logwarn("No ODOM for %.1fs. Check topic: %s", a_odom if a_odom else -1, self.odom_topic)

    def cb(self, img_msg, cloud_msg, odom_msg):
        t_cb0 = time.time()
        # ---- 1) 解码图像 ----
        try:
            if self._image_is_compressed:
                np_arr = np.frombuffer(img_msg.data, np.uint8)
                img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            else:
                img = self.bridge.imgmsg_to_cv2(img_msg, desired_encoding="bgr8")
        except CvBridgeError as e:
            rospy.logerr("cv_bridge error: %s", e); return
        if img is None:
            rospy.logerr("Decoded image is None"); return
        h, w = img.shape[:2]
        if self.cam_w <= 0 or self.cam_h <= 0:
            self.cam_w, self.cam_h = w, h

        self.msg_cnt_img += 1; self.last_stamp_img = img_msg.header.stamp
        rospy.logdebug("IMG #%d ts=%.3f size=%dx%d frame=%s",
                       self.msg_cnt_img, img_msg.header.stamp.to_sec(), w, h, img_msg.header.frame_id)

        # ---- 2) ODOM ----
        self.msg_cnt_odom += 1; self.last_stamp_odom = odom_msg.header.stamp
        q = odom_msg.pose.pose.orientation
        t = odom_msg.pose.pose.position
        R_w_i = R.from_quat([q.x, q.y, q.z, q.w]).as_matrix()
        t_w_i = np.array([t.x, t.y, t.z]).reshape(3,1)
        T_w_i = make_SE3(R_w_i, t_w_i)
        T_i_l = inv_SE3(self.T_l_i)
        T_w_c = T_w_i @ T_i_l @ self.T_l_c
        T_c_w = inv_SE3(T_w_c)
        if self.msg_cnt_odom <= 3:
            rospy.loginfo("ODOM #%d ts=%.3f frame=%s | T_w_c[0:3,3]=%s",
                          self.msg_cnt_odom, odom_msg.header.stamp.to_sec(), odom_msg.header.frame_id, T_w_c[:3,3])

        # ---- 3) 点云（PointXYZRGB）----
        pts_w, rgb = self.read_xyzrgb(cloud_msg)
        self.msg_cnt_cloud += 1; self.last_stamp_cloud = cloud_msg.header.stamp
        rospy.logdebug("CLOUD #%d ts=%.3f points=%d frame=%s",
                       self.msg_cnt_cloud, cloud_msg.header.stamp.to_sec(), pts_w.shape[0], cloud_msg.header.frame_id)
        if pts_w.size == 0:
            rospy.logwarn_throttle(2.0, "cloud has 0 points"); return

        # ---- 4) YOLO 检测 ----
        t0 = time.time()
        try:
            res = self.model.predict(source=img, imgsz=self.img_size, conf=self.conf_thres,
                                     iou=self.iou_thres, device=self.device, verbose=False)
            det = res[0]
            boxes = det.boxes.xyxy.cpu().numpy() if det.boxes is not None else np.zeros((0,4))
            cls   = det.boxes.cls.cpu().numpy().astype(int) if det.boxes is not None else np.zeros((0,), dtype=int)
            conf  = det.boxes.conf.cpu().numpy() if det.boxes is not None else np.zeros((0,), dtype=float)
            names = det.names if hasattr(det, "names") else {}
        except Exception as e:
            rospy.logerr("YOLO inference failed: %s\n%s", e, traceback.format_exc()); return
        t1 = time.time()
        rospy.loginfo_throttle(1.0, "YOLO: det=%d (%.1f ms)", boxes.shape[0], (t1-t0)*1000.0)

        # ---- 5) 世界->相机 投影 ----
        xyz_c = self.transform_pts(T_c_w, pts_w)            # (N,3)
        Z = xyz_c[:,2]
        depth_mask = (Z > self.min_depth) & (Z < self.max_depth)
        if self.fx <= 0 or self.fy <= 0:
            rospy.logerr("Invalid intrinsics fx/fy <= 0"); return
        u = self.fx * (xyz_c[:,0] / (Z + 1e-9)) + self.cx
        v = self.fy * (xyz_c[:,1] / (Z + 1e-9)) + self.cy
        in_img = (u >= 0) & (u < self.cam_w) & (v >= 0) & (v < self.cam_h)
        proj_valid = depth_mask & in_img
        n_in = int(np.count_nonzero(proj_valid))
        rospy.loginfo_throttle(1.0, "Project: pts=%d | in_img=%d (%.1f%%) | depth[%.2f,%.2f]",
                               pts_w.shape[0], n_in, 100.0*n_in/max(1,pts_w.shape[0]), self.min_depth, self.max_depth)

        # ---- 6) 染色统计 ----
        rgb_out = rgb.copy()
        colored_total = 0
        if boxes.shape[0] > 0 and n_in > 0:
            pal = self.palette(len(boxes))
            uv = np.stack([u[proj_valid], v[proj_valid]], axis=1)
            idx_valid = np.where(proj_valid)[0]
            for bi,(x1,y1,x2,y2) in enumerate(boxes):
                inside = (uv[:,0] >= x1) & (uv[:,0] <= x2) & (uv[:,1] >= y1) & (uv[:,1] <= y2)
                idx_pts = idx_valid[inside]
                rgb_out[idx_pts,:] = pal[bi]
                colored_total += idx_pts.size
            rospy.loginfo_throttle(1.0, "Colored pts in boxes: %d", colored_total)
        else:
            if boxes.shape[0] == 0:
                rospy.loginfo_throttle(1.0, "No detections this frame")
            if n_in == 0:
                rospy.logwarn_throttle(1.0, "No projected points fall in image — check extrinsics/frames/time-sync")

        # ---- 7) 发布图像/点云 & 保存调试图 ----
        img_mark = self.draw(img.copy(), boxes, cls, conf, names)
        try:
            self.pub_img.publish(self.bridge.cv2_to_imgmsg(img_mark, encoding="bgr8"))
        except Exception as e:
            rospy.logerr("publish image failed: %s", e)
        try:
            cv2.imwrite("/tmp/yolo_marked.jpg", img_mark)
        except Exception as e:
            rospy.logwarn("save /tmp/yolo_marked.jpg failed: %s", e)

        cloud_hl = self.build_cloud_rgb(pts_w, rgb_out, frame_id=cloud_msg.header.frame_id, stamp=cloud_msg.header.stamp)
        self.pub_cloud.publish(cloud_hl)

        rospy.loginfo_throttle(1.0, "End frame in %.1f ms", (time.time()-t_cb0)*1000.0)

    # ---------- utils ----------
    def read_xyzrgb(self, msg: PointCloud2):
        xyz = []; rgb = []
        try:
            # 若你的 cloud_registered 没有 r/g/b 字段，会抛异常
            for p in pc2.read_points(msg, field_names=("x","y","z","r","g","b"), skip_nans=True):
                xyz.append([p[0],p[1],p[2]])
                rgb.append([int(p[3]),int(p[4]),int(p[5])])
        except Exception as e:
            rospy.logerr_throttle(2.0, "read_points need fields x y z r g b — got error: %s", e)
            return np.empty((0,3)), np.empty((0,3), np.uint8)
        if len(xyz)==0:
            return np.empty((0,3)), np.empty((0,3), np.uint8)
        return np.asarray(xyz, np.float64), np.asarray(rgb, np.uint8)

    def transform_pts(self, T, pts_w):
        N = pts_w.shape[0]
        Pw = np.hstack([pts_w, np.ones((N,1))])
        Pc = (T @ Pw.T).T
        return Pc[:,:3].astype(np.float64)

    def build_cloud_rgb(self, xyz, rgb, frame_id, stamp):
        header = Header(); header.frame_id = frame_id; header.stamp = stamp
        fields = [
            pc2.PointField('x', 0,  pc2.PointField.FLOAT32, 1),
            pc2.PointField('y', 4,  pc2.PointField.FLOAT32, 1),
            pc2.PointField('z', 8,  pc2.PointField.FLOAT32, 1),
            pc2.PointField('r', 12, pc2.PointField.UINT8,   1),
            pc2.PointField('g', 13, pc2.PointField.UINT8,   1),
            pc2.PointField('b', 14, pc2.PointField.UINT8,   1),
        ]
        pts = []
        for i in range(xyz.shape[0]):
            pts.append([float(xyz[i,0]), float(xyz[i,1]), float(xyz[i,2]),
                        int(rgb[i,0]), int(rgb[i,1]), int(rgb[i,2])])
        return pc2.create_cloud(header, fields, pts)

    def draw(self, img, boxes, cls, conf, names):
        for i,(x1,y1,x2,y2) in enumerate(boxes):
            color = (int(255*(i%3==0)), int(255*(i%3==1)), int(255*(i%3==2)))
            cv2.rectangle(img, (int(x1),int(y1)), (int(x2),int(y2)), color, 2)
            if cls is not None and i < len(cls):
                name = names.get(cls[i], str(cls[i])) if isinstance(names, dict) else str(cls[i])
                s = conf[i] if conf is not None and i < len(conf) else 0.0
                label = f"{name} {s:.2f}"
                (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
                cv2.rectangle(img, (int(x1), int(y1)-th-4), (int(x1)+tw+4, int(y1)), color, -1)
                cv2.putText(img, label, (int(x1)+2, int(y1)-2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 1)
        return img

    def palette(self, n):
        rng = np.random.default_rng(2024)
        return rng.integers(low=64, high=255, size=(n,3), dtype=np.uint8)

def main():
    rospy.init_node("yolo_lidar_fusion", anonymous=False)
    Node()
    rospy.spin()

if __name__ == "__main__":
    main()

